Skip to content

Reduce Grouped MLP Fuser CPU Overhead - #3410

Open
zhongbozhu wants to merge 3 commits into
NVIDIA:mainfrom
zhongbozhu:opfuser_group_mlp_cpu_optim
Open

Reduce Grouped MLP Fuser CPU Overhead#3410
zhongbozhu wants to merge 3 commits into
NVIDIA:mainfrom
zhongbozhu:opfuser_group_mlp_cpu_optim

Conversation

@zhongbozhu

@zhongbozhu zhongbozhu commented Aug 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups.

Fixes # (issue)

#2897

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 20, 2026
@zhongbozhu
zhongbozhu marked this pull request as ready for review August 20, 2026 18:39
@zhongbozhu
zhongbozhu requested a review from timmoon10 as a code owner August 20, 2026 18:39
@zhongbozhu

zhongbozhu commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch L1

@greptile-apps

greptile-apps Bot commented Aug 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR reduces grouped-MLP fuser CPU overhead by caching operation plans across alternating checkpoint forward and recompute modes and by avoiding redundant current-device discovery.

  • Caches forward and backward operation lists by recipe type, backward boundary, and backward override.
  • Separates recipe-state reset from fused-plan selection.
  • Uses the grouped-MLP weight device directly for CUDA stream lookup.
  • Adds a focused test covering plan reuse as gradient requirements alternate.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/pytorch/ops/fuser.py Adds reusable fusion-plan caching while continuing to update runtime backward boundaries and recipe-owned basic-operation state.
transformer_engine/pytorch/ops/fused/grouped_mlp.py Selects the current CUDA stream using the grouped-MLP weight device index at both forward and backward launch sites.
tests/pytorch/test_fusible_ops.py Verifies that alternating no-grad and grad-enabled invocations construct two plans and subsequently restore the corresponding cached lists.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Fuser invocation] --> B[Determine backward boundary]
  B --> C{Recipe state changed?}
  C -->|Yes| D[Reset basic-operation recipe state]
  C -->|No| E[Build fusion cache key]
  D --> E
  E --> F{Cached plan exists?}
  F -->|Yes| G[Restore forward and backward lists]
  F -->|No| H[Apply fusion rules]
  H --> I[Cache constructed plan]
  G --> J[Execute selected plan]
  I --> J
Loading

Reviews (5): Last reviewed commit: "fix cutedsl wgrad crash" | Re-trigger Greptile

greptile-apps[bot]
greptile-apps Bot reviewed Aug 20, 2026
View reviewed changes
Comment thread transformer_engine/pytorch/ops/fuser.py
vthumbe1503
vthumbe1503 reviewed Aug 20, 2026
View reviewed changes
Comment thread transformer_engine/pytorch/ops/fuser.py
# backward boundary in the key, but pay construction cost only once for
# each configuration. Full recompute therefore builds at most one
# no-grad plan and one grad-enabled plan for a stable recipe.
fusion_params = (recipe_type, first_op_requiring_backward, backward_override)

@vthumbe1503 vthumbe1503 Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should use recipe object instead of recipe_type as fusion_param key And implement hash function for it. Since that governs the fusion and not just the recipe type

cc: @timmoon10

@vthumbe1503 vthumbe1503 Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@zhongbozhu commenting here so we dont miss out on our discussions related to this PR. The main issue we need to take care of here is the CPU overheads of the additional recipe object as a key in fusion params.

There are two options to deal with this

  1. We can cache the recipe representation so that the later iterations using the same recipe object dont have to deal with recomputing the representation/hash needed to be maintained in a fusion params dict. However it turns out recipe object can be recreated on every iteration (for eg, in Megatron).
  2. So @timmoon10 had the idea of exposing a representation of recipe object, call it options_key. Different objects will have the same options_key(). And so fuser will take the responsibility of caching the options_key instead of Recipe class itself taking that responsibility. And so if recipe is unchanged in terms of options_key, we wont face any new CPU overheads for the fusion_params_key computation

@tingyangk
tingyangk force-pushed the opfuser_group_mlp_cpu_optim branch 2 times, most recently from ff86f24 to 97850a4 Compare August 22, 2026 02:00
@zhongbozhu

zhongbozhu commented Aug 22, 2026

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch L1

@zhongbozhu

zhongbozhu commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

zhongbozhu and others added 3 commits August 27, 2026 13:38
@zhongbozhu @tingyangk
Reduce grouped MLP fuser CPU overhead
7e4386b 
Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups.

Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@zhongbozhu
resolve comments
4b18419 
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
@tingyangk @zhongbozhu
fix cutedsl wgrad crash
866c6f1 
Signed-off-by: tingyangk <tingyangk@nvidia.com>
@zhongbozhu
zhongbozhu force-pushed the opfuser_group_mlp_cpu_optim branch from 97850a4 to 866c6f1 Compare August 27, 2026 20:38
@zhongbozhu

zhongbozhu commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch L1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@greptile-apps greptile-apps[bot] greptile-apps[bot] left review comments
@vthumbe1503 vthumbe1503 vthumbe1503 left review comments
@timmoon10 timmoon10 Awaiting requested review from timmoon10 timmoon10 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@zhongbozhu @vthumbe1503 @tingyangk